package de.unibi.comet.hmm;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.NoSuchElementException;

import de.unibi.comet.util.LogSpace;


/** Implementation of StateAndMassDistribution interface using arrays. */
public class ArrayStateAndMassDistribution implements StateAndMassDistribution {
	private int states;
	private ArrayList<double[]> table;
	private int massBound;
	private boolean logarithmic;
	
	public ArrayStateAndMassDistribution(int states, int initialMaxMass, boolean logarithmic) {
		this.massBound=-1;
		this.logarithmic=logarithmic;
		this.states=states;
		table=new ArrayList<double[]>(initialMaxMass);
		growTo(initialMaxMass);
	}
	
	public ArrayStateAndMassDistribution(int states, int initialMaxMass, int massBound, boolean logarithmic) {
		if (initialMaxMass>massBound) throw new IllegalArgumentException();
		this.massBound=massBound;
		this.logarithmic=logarithmic;
		this.states=states;
		table=new ArrayList<double[]>(initialMaxMass);
		growTo(initialMaxMass);
	}

	private void growTo(int mass) {
		table.ensureCapacity(mass);
		for (int i=table.size(); i<=mass; ++i) {
			double[] d = new double[states];
			if (logarithmic) {
				for (int j=0; j<d.length; ++j) d[j]=Double.NEGATIVE_INFINITY;
			}
			table.add(d);
		}
	}
	
	/** Adds probability p to entry state,mass in the table.
	 *  Pleas note: if this is a logarithmic table, p is expected to
	 *  be given in log-space also.
	 */
	public void add(int state, int mass, double p) {
		if (mass>massBound) throw new IndexOutOfBoundsException();
		if (table.size()<=mass) growTo(mass);
		if (!logarithmic) {
			table.get(mass)[state]+=p;
		} else {
			double[] d = table.get(mass); 
			d[state]=LogSpace.logAdd(p,d[state]);
		}
	}

	public StateAndMassDistribution emptyClone() {
		return new ArrayStateAndMassDistribution(states, table.size()-1, massBound, logarithmic);
	}

	public double get(int state, int mass) {
		try {
			return table.get(mass)[state];
		} catch (ArrayIndexOutOfBoundsException e) {
			return 0.0;
		} catch (IndexOutOfBoundsException e) {
			return 0.0;
		}
	}

	private class MassDistributionIterator implements Iterator<Emission> {
		private int state;
		private int nextMass;
		private double nextMassProb;
		MassDistributionIterator(int state) {
			this.state=state;
			this.nextMass=0;
			this.nextMassProb=0.0;
			findNextMass(0);
		}
		private void findNextMass(int min) {
			for (int i=min; i<table.size(); ++i) {
				double p = table.get(i)[state];
				if ((!logarithmic && (p>0.0)) || (logarithmic && (p>Double.NEGATIVE_INFINITY))) {
					nextMass=i;
					nextMassProb=p;
					return;
				}
			}
			nextMass=-1;
		}
		public boolean hasNext() {
			return nextMass>=0;
		}
		public Emission next() {
			if (nextMass<0) throw new NoSuchElementException();
			Emission result = new Emission(nextMass, nextMassProb);
			findNextMass(nextMass+1);
			return result;
		}
		public void remove() { throw new UnsupportedOperationException(); }
	}
	
	public Iterator<Emission> getMassDistribution(int state) {
		return new MassDistributionIterator(state);
	}

	public void set(int state, int mass, double p) {
		if (table.size()<=mass) growTo(mass);
		table.get(mass)[state]=p;
	}

	public double getMassProbability(int mass) {
		double result = logarithmic?Double.NEGATIVE_INFINITY:0.0;
		for (double p : table.get(mass)) {
			if (!logarithmic) {
				result+=p;
			} else {
				result=LogSpace.logAdd(p, result);
			}
		}
		return result;
	}

	public double getStateProbability(int state) {
		double result = logarithmic?Double.NEGATIVE_INFINITY:0.0;
		for (double[] a : table) {
			if (!logarithmic) {
				result+=a[state];
			} else {
				result=LogSpace.logAdd(a[state], result);
			}
		}
		return result;
	}

	public double[] getStateDistribution() {
		double [] result = new double[states];
		if (logarithmic) {
			for (int i=0; i<result.length; ++i) {
				result[i]=Double.NEGATIVE_INFINITY;
			}
		}
		for (double[] a : table) {
			for (int i=0; i<a.length; ++i) {
				if (!logarithmic) {
					result[i]+=a[i];
				} else {
					result[i]=LogSpace.logAdd(result[i], a[i]);
				}
			}
		}
		return result;
	}
	
	public boolean isLogarithmic() {
		return logarithmic;
	}

	public int getMaxMass() {
		return table.size()-1;
	}
	
	public int getMassBound() {
		return massBound;
	}

	public int getStateCount() {
		return states;
	}

	@Override
	public String toString() {
		StringBuffer sb = new StringBuffer();
		for (int state=0; state<getStateCount(); ++state) {
			sb.append(String.format("%2d", state));
			for (int mass=0; mass<=getMaxMass(); ++mass) {
				// sb.append(String.format(" & %.3f", get(state, mass)));
				sb.append(String.format(" | %e", get(state, mass)));
			}
			sb.append(String.format(" || %e \n", getStateProbability(state)));
		}
		return sb.toString();
	}
	
}
